{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import pandas as pd\n",
    "import random\n",
    "\n",
    "file_path = r'/home/b8313/coding/py/high-frequency-trading/test_data/中国平安601318.SH.xlsx'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "outputs": [],
   "source": [
    "raw_file_data = pd.read_excel(file_path)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "outputs": [],
   "source": [
    "file_data = raw_file_data.head(20000)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "outputs": [],
   "source": [
    "action_choose = [0, 1, 2, 3, 4, 5, 6, 7, 8]"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "outputs": [],
   "source": [
    "init_positions_number = 0\n",
    "init_positions = file_data.iloc[0].收盘价\n",
    "\n",
    "init_current_line = 0"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "outputs": [
    {
     "data": {
      "text/plain": "<All keys matched successfully>"
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = torch.nn.Sequential(\n",
    "    torch.nn.Linear(11, 128),\n",
    "    torch.nn.ReLU(),\n",
    "    torch.nn.Linear(128, 9),\n",
    ")\n",
    "\n",
    "model.load_state_dict(torch.load('/home/b8313/Downloads/model_128_1667770743_20000.pth'))"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "outputs": [],
   "source": [
    "class StockEnv:\n",
    "    def __init__(self):\n",
    "        self.stock_data = file_data\n",
    "        self.over_line = len(self.stock_data)\n",
    "        self.action_choose = action_choose\n",
    "        self.current_line = init_current_line\n",
    "        self.positions_number = init_positions_number\n",
    "        self.positions_price = init_positions\n",
    "        self.rating = 500  # 赚钱的比例系数\n",
    "\n",
    "    def step(self, action):\n",
    "        this_time_stock = self.stock_data.iloc[self.current_line]\n",
    "        action2reward = -15\n",
    "        if action == 0:\n",
    "            # 买100手\n",
    "            self.positions_price = ((self.positions_price * self.positions_number) + 100 *\n",
    "                                    this_time_stock['价']) / (self.positions_number + 100)\n",
    "            self.positions_number = self.positions_number + 100\n",
    "            action2reward = -400\n",
    "        elif action == 1:\n",
    "            # 买50手\n",
    "            self.positions_price = ((self.positions_price * self.positions_number) + 50 *\n",
    "                                    this_time_stock['价']) / (self.positions_number + 50)\n",
    "            self.positions_number = self.positions_number + 50\n",
    "            action2reward = -200\n",
    "        elif action == 2:\n",
    "            # 买20手\n",
    "            self.positions_price = ((self.positions_price * self.positions_number) + 20 *\n",
    "                                    this_time_stock['价']) / (self.positions_number + 20)\n",
    "            self.positions_number = self.positions_number + 20\n",
    "            action2reward = -80\n",
    "        elif action == 3:\n",
    "            # 买10手\n",
    "            self.positions_price = ((self.positions_price * self.positions_number) + 10 *\n",
    "                                    this_time_stock['价']) / (self.positions_number + 10)\n",
    "            self.positions_number = self.positions_number + 10\n",
    "            action2reward = -20\n",
    "        elif action == 4:\n",
    "            pass\n",
    "        elif action == 5:\n",
    "            # 卖10手\n",
    "            if self.positions_number < 10:\n",
    "                action2reward = -100000\n",
    "            else:\n",
    "                action2reward = self.rating * (this_time_stock['价'] - self.positions_price) * 10 - 10  # 赚到0元也要扣10reward\n",
    "                self.positions_number = self.positions_number - 10\n",
    "        elif action == 6:\n",
    "            # 卖20手\n",
    "            if self.positions_number < 20:\n",
    "                action2reward = -200000\n",
    "            else:\n",
    "                action2reward = self.rating * (this_time_stock['价'] - self.positions_price) * 20 - 10\n",
    "                self.positions_number = self.positions_number - 20\n",
    "        elif action == 7:\n",
    "            # 卖50手\n",
    "            if self.positions_number < 50:\n",
    "                action2reward = -500000\n",
    "            else:\n",
    "                action2reward = self.rating * (this_time_stock['价'] - self.positions_price) * 50 - 10\n",
    "                self.positions_number = self.positions_number - 50\n",
    "        elif action == 8:\n",
    "            # 卖100手\n",
    "            if self.positions_number < 100:\n",
    "                action2reward = -1000000\n",
    "            else:\n",
    "                action2reward = self.rating * (this_time_stock['价'] - self.positions_price) * 100 - 10\n",
    "                self.positions_number = self.positions_number - 100\n",
    "        else:\n",
    "            pass\n",
    "\n",
    "        self.current_line += 1\n",
    "        # print(self.current_line)\n",
    "        current_stock = self.stock_data.iloc[self.current_line]\n",
    "        next_state = (current_stock['收盘价'], current_stock['开盘价'], current_stock['最高价'],\n",
    "                      current_stock['最低价'], current_stock['成交额（元）'], current_stock['成交量（手）'],\n",
    "                      current_stock['当前额'], current_stock['当前量'], current_stock['价'],\n",
    "                      self.positions_number, self.positions_price)\n",
    "        over = self.current_line == self.over_line - 1\n",
    "        return next_state, action2reward, over, None\n",
    "\n",
    "    # 返回null只是为了对比\n",
    "\n",
    "    def get_current_state(self):\n",
    "        current_stock = self.stock_data.iloc[self.current_line]\n",
    "        state = (current_stock['收盘价'], current_stock['开盘价'], current_stock['最高价'],\n",
    "                 current_stock['最低价'], current_stock['成交额（元）'], current_stock['成交量（手）'],\n",
    "                 current_stock['当前额'], current_stock['当前量'], current_stock['价'],\n",
    "                 self.positions_number, self.positions_price)\n",
    "        return state\n",
    "\n",
    "    def get_action(self, state):\n",
    "        if random.random() < 0.01:\n",
    "            # print(\"随机了一个动作\")\n",
    "            return random.choice(self.action_choose)\n",
    "\n",
    "        #走神经网络,得到一个动作\n",
    "        # print('走神经网络,得到一个动作')\n",
    "        reshape_state = torch.FloatTensor(state).reshape(1, 11)\n",
    "\n",
    "        return model(reshape_state).argmax().item()\n",
    "\n",
    "    def reset(self):\n",
    "        self.current_line = init_current_line\n",
    "        self.positions_number = init_positions_number\n",
    "        self.positions_price = init_positions\n",
    "\n",
    "        return self.get_current_state()\n",
    "\n",
    "\n",
    "env = StockEnv()"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 4, reward -15\n",
      "action 1, reward -200\n",
      "action 0, reward -400\n",
      "action 7, reward 240.0000000001279\n",
      "action 7, reward 240.0000000001279\n",
      "action 5, reward 90.00000000001563\n",
      "action 5, reward -10.0\n",
      "action 7, reward -500000\n",
      "action 7, reward -500000\n",
      "action 7, reward -500000\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 7, reward 181.86046511627808\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 7, reward -728.287526427197\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 7, reward 6328.061772716941\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n",
      "action 0, reward -400\n"
     ]
    }
   ],
   "source": [
    "# test_state=env.get_current_state()\n",
    "# reshape_state = torch.FloatTensor(test_state).reshape(1, 11)\n",
    "\n",
    "model.eval()\n",
    "# env.reset()\n",
    "test_state = env.get_current_state()\n",
    "reshape_state = torch.FloatTensor(test_state).reshape(1, 11)\n",
    "for turn in range(500):\n",
    "    action = model(reshape_state).argmax().item()\n",
    "    # print(action)\n",
    "    next_state, action2reward, over, _ = env.step(action)\n",
    "    print(\"action {}, reward {}\".format(action,action2reward))\n",
    "    reshape_state = torch.FloatTensor(next_state).reshape(1, 11)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
